import pandas as pd
from sklearn.model_selection import StratifiedShuffleSplit, train_test_split

def extract_binary_label(df: pd.DataFrame, label_col: str = "cogimp_label") -> pd.Series:
    """
    从 DataFrame 提取二分类标签    -1 → 0 (不生病)    1~9 → 1 (生病)
    :param df: 完整数据集
    :param label_col: 标签列
    :return: 二分类标签
    :raises: ValueError: 如果 DataFrame 中没有标签列
    """
    if label_col not in df.columns:
        raise ValueError(f"列 {label_col} 不在 DataFrame 中！")
    
    y = df[label_col].copy()

    # 映射规则：-1 → 0，其余正数 → 1
    y_binary = (y != -1).astype(int)

    return y_binary

def extract_regression_label(df: pd.DataFrame, label_suffix: str = "cogimpt_label") -> pd.DataFrame:
    """
    从 DataFrame 提取回归标签（数值型标签）
    
    :param df: 数据集
    :param label_suffix: 标签列后缀（默认提取所有以该后缀结尾的列）
    :return: 回归标签 DataFrame（支持多输出回归）
    :raises: ValueError: 如果没有找到标签列或存在非数值型标签
    """
    # 找出所有以 label_suffix 结尾的列
    y_cols = [col for col in df.columns if col.endswith(label_suffix)]
    if not y_cols:
        raise ValueError(f"没有找到以 {label_suffix} 结尾的标签列！")

    y = df[y_cols].copy()

    # 检查每一列是否是数值型
    non_numeric = [col for col in y_cols if not pd.api.types.is_numeric_dtype(df[col])]
    if non_numeric:
        raise ValueError(f"标签列 {non_numeric} 不是数值型！实际类型: {df[non_numeric].dtypes}")

    return y

def split_data(df: pd.DataFrame, id_col: str="subject_id", label_col: str="cogimp_label", test_size=0.2, random_state=42):
    """
    按 subject_id 分组并分层划分 train/test
    :param df: 数据集
    :param id_col: 样本 id 列
    :param label_col: 标签列
    :param test_size: 测试集比例
    :param random_state: 随机种子
    :return: 训练集和测试集
    :raises: ValueError: 如果 DataFrame 中没有 id_col 列
    """
    if id_col not in df.columns:
        raise ValueError(f"DataFrame 中没有 {id_col} 列！")

    # 统计每个 subject 的主要标签（取第一个，或者用多数标签）
    subject_labels = df.groupby(id_col)[label_col].agg(lambda x: x.mode()[0]).reset_index()

    sss = StratifiedShuffleSplit(n_splits=1, test_size=test_size, random_state=random_state)
    train_idx, test_idx = next(sss.split(subject_labels[id_col], subject_labels[label_col]))

    train_ids = subject_labels.iloc[train_idx][id_col].tolist()
    test_ids = subject_labels.iloc[test_idx][id_col].tolist()

    train_df = df[df[id_col].isin(train_ids)].reset_index(drop=True)
    test_df = df[df[id_col].isin(test_ids)].reset_index(drop=True)

    return train_df, test_df

import torch
torch.set_num_threads(1)

def generate_dummy_batch(dataset, batch_size=1, device='cuda'):
    """
    从 Dataset 自动生成合法的 dummy batch
    :param dataset: ADHDDataset
    :param batch_size: batch 大小
    :param device: 'cpu' 或 'cuda'
    :return: x_cat, x_cont, y 都是 batch 形式
    """
    x_cat_list, x_cont_list, y_list = [], [], []

    for _ in range(batch_size):
        # 生成合法的类别值
        x_cat = torch.tensor(
            [torch.randint(0, dataset.cat_encoders[col].classes_.size, (1,)).item()
             for col in dataset.cat_cols],
            dtype=torch.long
        )

        # 生成随机连续值（使用 dataset 的 scaler 范围可选）
        x_cont = torch.randn(len(dataset.num_cols), dtype=torch.float)

        # 生成合法标签
        y = torch.randint(0, len(dataset.labels.unique()), (1,), dtype=torch.long).squeeze(0)

        x_cat_list.append(x_cat.unsqueeze(0))
        x_cont_list.append(x_cont.unsqueeze(0))
        y_list.append(y.unsqueeze(0))

    # 拼成 batch
    x_cat_batch = torch.cat(x_cat_list, dim=0).to(device)
    x_cont_batch = torch.cat(x_cont_list, dim=0).to(device)
    y_batch = torch.cat(y_list, dim=0).to(device)

    return x_cat_batch, x_cont_batch, y_batch

